#version 460

#extension GL_EXT_scalar_block_layout : enable
#extension GL_EXT_ray_tracing : require
#extension GL_EXT_ray_query : enable
#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_nonuniform_qualifier : require
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
#extension GL_EXT_buffer_reference2 : require

#define RAY_TRACING
#include "../Brdf.glsl"
#include "DdgiCommon.glsl"

layout (binding = 0) uniform accelerationStructureEXT topLevelAS;

#include "../RayQuery.glsl"
#define RAY_THROUGHPUT
#define SAMPLE_SKY_LIGHT
#include "../Lighting.glsl"

layout (buffer_reference, scalar) buffer Vertices {Vertex v[]; };

layout (buffer_reference, scalar) buffer Indices {ivec3 i[]; };

layout (set = 0, binding = 3, scalar) uniform DDGIUBO
{
    DDGIUniforms data;
} ddgiUbo;

layout (set = 0, binding = 4, scalar) buffer ObjectDescBuffer { ObjectDesc data[]; } objectDescs;

layout (set = 0, binding = 5, scalar) buffer MaterialBuffer { Material data[]; } materials;

layout (set = 0, binding = 6) uniform PerFrameUBO
{
    mat4 viewInverse;
    mat4 projInverse;
    mat4 projViewInverse;
    mat4 prevProjView;
    mat4 projView;
    vec4 cameraPosition;
    vec4 currentPrevJitter;
    Light light;
} perFrameUBO;

layout (set = 0, binding = 7) uniform samplerCube cubemapImage;

layout (set = 1, binding = 0) uniform sampler2D irradianceImage;

layout (set = 1, binding = 1) uniform sampler2D depthImage;


layout(push_constant) uniform PushConstants
{
    mat4 randomOrientation;
    uint frame;
    uint infiniteBounces;
    float giIntensity;
} pushConstants;

layout(location = 0) rayPayloadInEXT DDGIPayload hitPayload;

hitAttributeEXT vec2 hitAttribs;


vec3 FresnelSchlickRoughness(float cosTheta, vec3 F0, float roughness)
{
    return F0 + (max(vec3(1.0 - roughness), F0) - F0) * pow(max(1.0 - cosTheta, 0.0), 5.0);
}


vec3 IndirectLighting(vec3 Wo, vec3 N, vec3 P, vec3 F0, vec3 diffuseColor, float roughness, float metallic)
{
    vec3 F = FresnelSchlickRoughness(max(dot(N, Wo), 0.0), F0, roughness);

    vec3 kS = F;
    vec3 kD = 1.0 - kS;
    kD *= 1.0 - metallic;

    return pushConstants.giIntensity * kD * diffuseColor * SampleIrradiance(ddgiUbo.data, P, N, Wo, irradianceImage, depthImage);
}


void main()
{
    const ObjectDesc objectDesc = objectDescs.data[gl_InstanceCustomIndexEXT];
    const Material material = materials.data[gl_InstanceCustomIndexEXT];

    Indices indices = Indices(objectDesc.indexAddress);
    Vertices vertices = Vertices(objectDesc.vertexAddress);

    // Indices of the triangle
    ivec3 index = indices.i[gl_PrimitiveID];

    // Vertex of the triangle
    Vertex v0 = vertices.v[index.x];
    Vertex v1 = vertices.v[index.y];
    Vertex v2 = vertices.v[index.z];

    const vec3 barycentrics = vec3(1.0 - hitAttribs.x - hitAttribs.y, hitAttribs.x, hitAttribs.y);

    const vec3 position = vec3(gl_ObjectToWorldEXT * vec4(v0.position * barycentrics.x + v1.position * barycentrics.y + v2.position * barycentrics.z, 1.0));

    const vec3 N = normalize(vec3((v0.normal * barycentrics.x + v1.normal * barycentrics.y + v2.normal * barycentrics.z) * gl_WorldToObjectEXT));

    const vec3 Wo = -gl_WorldRayDirectionEXT;
    const vec3 R = reflect(-Wo, N);

    const vec3 albedo = material.albedo.rgb;
    const float roughness = material.roughnessMetallic.r;
    const float metallic = material.roughnessMetallic.g;

    const vec3 F0 = mix(vec3(0.04f), albedo, metallic);
    const vec3 diffuseColor = mix(albedo * (vec3(1.0f) - F0), vec3(0.0f), metallic);

    vec3 Lo = vec3(0.0f);

    Lo += DirectLighting(perFrameUBO.light, Wo, N, position, F0, diffuseColor, roughness, hitPayload.T, NextVec2(hitPayload.rng), cubemapImage);

    if (pushConstants.infiniteBounces == 1)
        Lo += IndirectLighting(Wo, N, position, F0, diffuseColor, roughness, metallic);

    hitPayload.L = Lo;
    hitPayload.hitDistance = gl_RayTminEXT + gl_HitTEXT;
}
